flock-core 0.1.1__py3-none-any.whl

flock/core/flock.py

@@ -0,0 +1,208 @@
+# 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

flock/core/handoff/handoff_base.py

@@ -0,0 +1,12 @@
+from typing import Any
+
+from flock.core.agent import Agent
+from flock.core.context import FlockContext
+
+
+class HandoffBase:
+    """Base class for handoff implementations."""
+
+    next_agent: str | Agent
+    input: dict[str, Any]
+    context: FlockContext

flock/core/logging/__init__.py

@@ -0,0 +1,18 @@
+"""Flock logging system with Rich integration and Temporal compatibility."""
+
+from flock.core.logging.error_handler import error_handler
+from flock.core.logging.formatters import PerformanceFormatter, StructuredFormatter
+from flock.core.logging.handlers import live_update_handler, performance_handler
+from flock.core.logging.logger import flock_logger
+
+# Install the Rich error handler by default
+error_handler.install()
+
+__all__ = [
+    "PerformanceFormatter",
+    "StructuredFormatter",
+    "error_handler",
+    "flock_logger",
+    "live_update_handler",
+    "performance_handler",
+]

flock/core/logging/error_handler.py

@@ -0,0 +1,84 @@
+import sys
+from types import TracebackType
+
+from temporalio import workflow
+
+with workflow.unsafe.imports_passed_through():
+    from rich.console import Console
+    from rich.panel import Panel
+    from rich.syntax import Syntax
+    from rich.traceback import Traceback
+
+from flock.core.logging.logger import flock_logger
+
+
+class ErrorHandler:
+    """Handles error formatting and display using Rich."""
+
+    def __init__(self, console: Console | None = None):
+        self.console = console or Console()
+
+    def format_exception(
+        self,
+        exc_type: type[BaseException],
+        exc_value: BaseException,
+        exc_tb: TracebackType,
+        *,
+        show_locals: bool = True,
+    ) -> Panel:
+        """Format an exception with Rich styling."""
+        # Create a Rich traceback
+        rich_tb = Traceback.from_exception(
+            exc_type,
+            exc_value,
+            exc_tb,
+            show_locals=show_locals,
+        )
+
+        # If there's source code available, syntax highlight it
+        if exc_tb and exc_tb.tb_frame.f_code.co_filename != "<string>":
+            try:
+                with open(exc_tb.tb_frame.f_code.co_filename) as f:
+                    source = f.read()
+                syntax = Syntax(
+                    source,
+                    "python",
+                    line_numbers=True,
+                    highlight_lines={exc_tb.tb_lineno},
+                )
+            except:
+                syntax = None
+        else:
+            syntax = None
+
+        # Create a panel with the traceback
+        return Panel(
+            rich_tb,
+            title=f"[red]{exc_type.__name__}[/]: {exc_value!s}",
+            border_style="red",
+            padding=(1, 2),
+        )
+
+    def handle_exception(
+        self,
+        exc_type: type[BaseException],
+        exc_value: BaseException,
+        exc_tb: TracebackType,
+    ) -> None:
+        """Handle an exception by formatting and displaying it."""
+        panel = self.format_exception(exc_type, exc_value, exc_tb)
+        self.console.print(panel)
+        # Also log the error through our logger
+        flock_logger.error(
+            f"Exception occurred: {exc_type.__name__}: {exc_value!s}",
+            error_type=exc_type.__name__,
+            error_details=str(exc_value),
+        )
+
+    def install(self) -> None:
+        """Install this error handler as the default exception handler."""
+        sys.excepthook = self.handle_exception
+
+
+# Global error handler instance
+error_handler = ErrorHandler()

flock/core/logging/formatters.py

@@ -0,0 +1,122 @@
+from typing import Any
+
+from devtools import pprint
+from temporalio import workflow
+
+with workflow.unsafe.imports_passed_through():
+    from rich.console import Console
+    from rich.panel import Panel
+    from rich.progress import BarColumn, Progress, SpinnerColumn, TextColumn, TimeElapsedColumn
+    from rich.table import Table
+
+
+class StructuredFormatter:
+    """Formats structured data for Rich output."""
+
+    @staticmethod
+    def create_status_table(data: dict[str, Any]) -> Table:
+        """Create a Rich table for displaying status information."""
+        table = Table(show_header=True, header_style="bold magenta")
+        table.add_column("Key")
+        table.add_column("Value")
+
+        for key, value in data.items():
+            table.add_row(str(key), str(value))
+
+        return table
+
+    @staticmethod
+    def create_progress_bar(description: str = "Progress") -> Progress:
+        """Create a Rich progress bar with time tracking."""
+        return Progress(
+            SpinnerColumn(),
+            TextColumn("[progress.description]{task.description}"),
+            BarColumn(),
+            TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
+            TimeElapsedColumn(),
+        )
+
+    @staticmethod
+    def create_workflow_panel(
+        workflow_id: str, status: str, details: dict[str, Any], title: str = "Workflow Status"
+    ) -> Panel:
+        """Create a Rich panel for workflow information."""
+        table = StructuredFormatter.create_status_table({"Workflow ID": workflow_id, "Status": status, **details})
+        return Panel(table, title=title, border_style="blue")
+
+    @staticmethod
+    def create_activity_panel(
+        activity_id: str, name: str, status: str, details: dict[str, Any], title: str = "Activity Status"
+    ) -> Panel:
+        """Create a Rich panel for activity information."""
+        table = StructuredFormatter.create_status_table(
+            {"Activity ID": activity_id, "Name": name, "Status": status, **details}
+        )
+        return Panel(table, title=title, border_style="magenta")
+
+
+class PerformanceFormatter:
+    """Formats performance metrics for Rich output."""
+
+    @staticmethod
+    def create_timing_tree(timings: dict[str, float]) -> str:
+        """Create a tree-like structure showing execution timings."""
+        if not timings:
+            return "No timings recorded"
+
+        result = "Performance Metrics:\n"
+        for operation, duration in timings.items():
+            result += f"├── {operation}: {duration:.3f}s\n"
+        return result[:-1]  # Remove trailing newline
+
+    @staticmethod
+    def create_performance_panel(metrics: dict[str, Any], title: str = "Performance Metrics") -> Panel:
+        """Create a panel showing performance metrics."""
+        table = StructuredFormatter.create_status_table(metrics)
+        return Panel(table, title=title, border_style="cyan")
+
+
+class AgentResultFormatter:
+    """Formats agent results in a beautiful Rich table."""
+
+    @staticmethod
+    def format_result(result: dict[str, Any], agent_name: str) -> Panel:
+        """Format an agent's result as a Rich panel containing a table."""
+        # Create a table with a nice header
+        table = Table(
+            show_header=True,
+            header_style="bold green",
+            title=f"Agent Results: {agent_name}",
+            title_style="bold blue",
+            border_style="bright_blue",
+        )
+        table.add_column("Output", style="cyan")
+        table.add_column("Value", style="green")
+
+        # Add each result to the table
+        for key, value in result.items():
+            # Format multi-line values (like blog headers) nicely
+            if isinstance(value, (list, tuple)) or (isinstance(value, str) and "\n" in value):
+                formatted_value = "\n".join(value) if isinstance(value, (list, tuple)) else value
+                # Add some padding for multi-line values
+                table.add_row(key, f"\n{formatted_value}\n")
+            else:
+                table.add_row(key, str(value))
+
+        pprint(result)
+
+        # Wrap the table in a panel for a nice border
+        return Panel(
+            table,
+            title="🎯 Agent Output",
+            title_align="left",
+            border_style="blue",
+            padding=(1, 2),
+        )
+
+    @staticmethod
+    def print_result(result: dict[str, Any], agent_name: str, console: Console | None = None) -> None:
+        """Print an agent's result using Rich formatting."""
+        console = console or Console()
+        panel = AgentResultFormatter.format_result(result, agent_name)
+        console.print(panel)

flock/core/logging/handlers.py

@@ -0,0 +1,117 @@
+import time
+
+from temporalio import workflow
+
+with workflow.unsafe.imports_passed_through():
+    from rich.console import Console
+    from rich.live import Live
+
+from collections.abc import Callable, Generator
+from contextlib import contextmanager
+from functools import wraps
+from typing import Any, TypeVar
+
+from flock.core.logging.formatters import PerformanceFormatter, StructuredFormatter
+
+T = TypeVar("T")
+
+
+class PerformanceHandler:
+    """Handles performance tracking and reporting."""
+
+    def __init__(self, console: Console | None = None):
+        self.console = console or Console()
+        self.timings: dict[str, float] = {}
+        self.enabled = False
+
+    def enable(self):
+        """Enable performance tracking."""
+        self.enabled = True
+
+    def disable(self):
+        """Disable performance tracking."""
+        self.enabled = False
+
+    def _get_time(self) -> float:
+        """Get current time in a workflow-safe way."""
+        try:
+            # Try to get workflow time first
+            return workflow.now().timestamp()
+        except workflow._NotInWorkflowEventLoopError:
+            # Fall back to system time if not in workflow
+            return time.time()
+
+    @contextmanager
+    def track_time(self, operation_name: str) -> Generator[None, None, None]:
+        """Context manager for tracking operation execution time."""
+        if not self.enabled:
+            yield
+            return
+
+        start_time = self._get_time()
+        try:
+            yield
+        finally:
+            end_time = self._get_time()
+            duration = end_time - start_time
+            self.timings[operation_name] = duration
+
+    def track_operation(self, operation_name: str) -> Callable[[Callable[..., T]], Callable[..., T]]:
+        """Decorator for tracking function execution time."""
+
+        def decorator(func: Callable[..., T]) -> Callable[..., T]:
+            @wraps(func)
+            def wrapper(*args: Any, **kwargs: Any) -> T:
+                with self.track_time(operation_name):
+                    return func(*args, **kwargs)
+
+            return wrapper
+
+        return decorator
+
+    def display_timings(self) -> None:
+        """Display all tracked timings in a tree format."""
+        if self.timings:
+            tree = PerformanceFormatter.create_timing_tree(self.timings)
+            self.console.print(tree)
+
+    def clear_timings(self) -> None:
+        """Clear all tracked timings."""
+        self.timings.clear()
+
+
+class LiveUpdateHandler:
+    """Handles live updates for long-running operations."""
+
+    def __init__(self, console: Console | None = None):
+        self.console = console or Console()
+
+    @contextmanager
+    def progress_tracker(self, description: str = "Progress") -> Generator[Any, None, None]:
+        """Context manager for tracking progress with a live display."""
+        progress = StructuredFormatter.create_progress_bar(description)
+        with Live(progress, console=self.console, refresh_per_second=4) as live:
+            task = progress.add_task(description, total=100)
+            try:
+                yield lambda p: progress.update(task, completed=p)
+            finally:
+                progress.update(task, completed=100)
+
+    def update_workflow_status(
+        self, workflow_id: str, status: str, details: dict[str, Any], refresh_per_second: int = 1
+    ) -> Live:
+        """Create a live updating workflow status panel."""
+        panel = StructuredFormatter.create_workflow_panel(workflow_id, status, details)
+        return Live(panel, console=self.console, refresh_per_second=refresh_per_second)
+
+    def update_activity_status(
+        self, activity_id: str, name: str, status: str, details: dict[str, Any], refresh_per_second: int = 1
+    ) -> Live:
+        """Create a live updating activity status panel."""
+        panel = StructuredFormatter.create_activity_panel(activity_id, name, status, details)
+        return Live(panel, console=self.console, refresh_per_second=refresh_per_second)
+
+
+# Global instances for convenience
+performance_handler = PerformanceHandler()
+live_update_handler = LiveUpdateHandler()

flock/core/logging/logger.py

@@ -0,0 +1,107 @@
+from typing import Any
+
+from temporalio import workflow
+
+from flock.core.logging.formatters import AgentResultFormatter
+
+with workflow.unsafe.imports_passed_through():
+    from rich.console import Console
+    from rich.theme import Theme
+
+# Custom theme for different log levels
+THEME = Theme({
+    "info": "cyan",
+    "warning": "yellow",
+    "error": "red",
+    "debug": "grey50",
+    "success": "green",
+    "workflow": "blue",
+    "activity": "magenta",
+})
+
+
+class LogLevel:
+    NONE = 0  # No logging except errors
+    MINIMAL = 1  # Only agent outputs
+    BASIC = 2  # Agent outputs + basic workflow info
+    VERBOSE = 3  # Everything including debug info
+
+
+class FlockLogger:
+    """Custom logger for Flock that integrates with Rich and respects Temporal constraints."""
+
+    def __init__(self, name: str = "flock", console: Console | None = None, level: int = LogLevel.MINIMAL):
+        self.name = name
+        self.console = console or Console(theme=THEME)
+        self._workflow_id: str | None = None
+        self._activity_id: str | None = None
+        self.level = level
+
+    def set_level(self, level: int) -> None:
+        """Set the logging level."""
+        self.level = level
+
+    def set_context(self, workflow_id: str | None = None, activity_id: str | None = None) -> None:
+        """Set the current workflow and activity context."""
+        self._workflow_id = workflow_id
+        self._activity_id = activity_id
+
+    def _format_message(self, level: str, message: str, **kwargs: Any) -> str:
+        """Format log message with context but without timestamps (for Temporal compatibility)."""
+        context_parts = []
+        if self._workflow_id:
+            context_parts.append(f"[workflow]workflow={self._workflow_id}[/]")
+        if self._activity_id:
+            context_parts.append(f"[activity]activity={self._activity_id}[/]")
+
+        # Add any additional context from kwargs
+        for key, value in kwargs.items():
+            if key not in ("workflow_id", "activity_id"):
+                context_parts.append(f"{key}={value}")
+
+        context_str = " ".join(context_parts)
+        return f"[{level}]{level.upper()}[/] {message} {context_str if context_str else ''}"
+
+    def info(self, message: str, **kwargs: Any) -> None:
+        """Log an info message."""
+        if self.level >= LogLevel.BASIC:
+            self.console.print(self._format_message("info", message, **kwargs))
+
+    def warning(self, message: str, **kwargs: Any) -> None:
+        """Log a warning message."""
+        if self.level >= LogLevel.BASIC:
+            self.console.print(self._format_message("warning", message, **kwargs))
+
+    def error(self, message: str, **kwargs: Any) -> None:
+        """Log an error message."""
+        # Errors are always logged regardless of level
+        self.console.print(self._format_message("error", message, **kwargs))
+
+    def debug(self, message: str, **kwargs: Any) -> None:
+        """Log a debug message."""
+        if self.level >= LogLevel.VERBOSE:
+            self.console.print(self._format_message("debug", message, **kwargs))
+
+    def success(self, message: str, **kwargs: Any) -> None:
+        """Log a success message."""
+        if self.level >= LogLevel.BASIC:
+            self.console.print(self._format_message("success", message, **kwargs))
+
+    def workflow_event(self, message: str, **kwargs: Any) -> None:
+        """Log a workflow-specific event."""
+        if self.level >= LogLevel.BASIC:
+            self.console.print(self._format_message("workflow", message, **kwargs))
+
+    def activity_event(self, message: str, **kwargs: Any) -> None:
+        """Log an activity-specific event."""
+        if self.level >= LogLevel.BASIC:
+            self.console.print(self._format_message("activity", message, **kwargs))
+
+    def result(self, message: dict, agent_name, **kwargs: Any) -> None:
+        """Log a result message."""
+        if self.level >= LogLevel.MINIMAL:
+            AgentResultFormatter.print_result(message, agent_name, self.console)
+
+
+# Global logger instance with minimal logging by default
+flock_logger = FlockLogger(level=LogLevel.MINIMAL)