flock-core 0.4.0b26__py3-none-any.whl → 0.4.0b28__py3-none-any.whl

flock/workflow/agent_execution_activity.py

@@ -0,0 +1,228 @@
+"""Defines granular Temporal activities for executing a single agent
+and determining the next agent in a Flock workflow.
+"""
+
+from collections.abc import Callable
+
+from opentelemetry import trace
+from temporalio import activity
+
+# Third-party imports only within activity functions if needed, or pass context
+# For core flock types, import directly
+from flock.core.context.context import FlockContext
+from flock.core.context.context_vars import FLOCK_MODEL
+from flock.core.flock_agent import FlockAgent  # Import concrete class if needed
+from flock.core.flock_registry import get_registry
+from flock.core.flock_router import HandOffRequest
+from flock.core.logging.logging import get_logger
+from flock.core.util.input_resolver import resolve_inputs
+
+logger = get_logger("agent_activity")  # Using a distinct logger category
+tracer = trace.get_tracer(__name__)
+registry = get_registry()  # Get registry instance once
+
+
+@activity.defn
+async def execute_single_agent(agent_name: str, context: FlockContext) -> dict:
+    """Executes a single specified agent and returns its result.
+
+    Args:
+        agent_name: The name of the agent to execute.
+        context: The current FlockContext (passed from the workflow).
+
+    Returns:
+        The raw result dictionary from the agent's execution.
+
+    Raises:
+        ValueError: If the agent is not found in the registry.
+        Exception: Propagates exceptions from agent execution for Temporal retries.
+    """
+    with tracer.start_as_current_span("execute_single_agent") as span:
+        span.set_attribute("agent.name", agent_name)
+        logger.info("Executing single agent", agent=agent_name)
+
+        agent = registry.get_agent(agent_name)
+        if not agent:
+            logger.error("Agent not found in registry", agent=agent_name)
+            # Raise error for Temporal to potentially retry/fail the activity
+            raise ValueError(f"Agent '{agent_name}' not found in registry.")
+
+        # Set agent's context reference (transient, for this execution)
+        agent.context = context
+
+        # Ensure model is set (using context value if needed)
+        # Consider if this should be done once when agent is added or workflow starts
+        if agent.model is None:
+            agent_model = context.get_variable(FLOCK_MODEL)
+            if agent_model:
+                agent.set_model(agent_model)
+                logger.debug(
+                    f"Set model for agent '{agent_name}' from context: {agent_model}"
+                )
+
+        # Resolve agent-specific callables if necessary
+        # This might be better handled in the workflow before the loop starts
+        # or when agents are initially loaded. Assuming it's handled elsewhere for now.
+        # agent.resolve_callables(context=context)
+
+        # Resolve inputs for this specific agent run
+        previous_agent_name = (
+            context.get_last_agent_name()
+        )  # Relies on context method
+        logger.debug(
+            f"Resolving inputs for {agent_name} with previous agent {previous_agent_name}"
+        )
+        agent_inputs = resolve_inputs(agent.input, context, previous_agent_name)
+        span.add_event(
+            "resolved inputs", attributes={"inputs": str(agent_inputs)}
+        )
+
+        try:
+            # Execute just this agent
+            result = await agent.run_async(agent_inputs)
+            # Avoid logging potentially large results directly to span attributes
+            result_str = str(result)
+            span.set_attribute("result.type", type(result).__name__)
+            span.set_attribute(
+                "result.preview",
+                result_str[:500] + ("..." if len(result_str) > 500 else ""),
+            )
+            logger.info("Single agent execution completed", agent=agent_name)
+            return result
+        except Exception as e:
+            logger.error(
+                "Single agent execution failed",
+                agent=agent_name,
+                error=str(e),
+                exc_info=True,
+            )
+            span.record_exception(e)
+            # Re-raise the exception for Temporal to handle based on retry policy
+            raise
+
+
+@activity.defn
+async def determine_next_agent(
+    current_agent_name: str, result: dict, context: FlockContext
+) -> dict | None:
+    """Determines the next agent using the current agent's handoff router.
+
+    Args:
+        current_agent_name: The name of the agent that just ran.
+        result: The result produced by the current agent.
+        context: The current FlockContext.
+
+    Returns:
+        A dictionary representing the HandOffRequest (serialized via model_dump),
+        or None if no handoff occurs or router doesn't specify a next agent.
+
+    Raises:
+        ValueError: If the current agent cannot be found.
+        Exception: Propagates exceptions from router execution for Temporal retries.
+    """
+    with tracer.start_as_current_span("determine_next_agent") as span:
+        span.set_attribute("agent.name", current_agent_name)
+        logger.info("Determining next agent after", agent=current_agent_name)
+
+        agent = registry.get_agent(current_agent_name)
+        if not agent:
+            logger.error(
+                "Agent not found for routing", agent=current_agent_name
+            )
+            raise ValueError(
+                f"Agent '{current_agent_name}' not found for routing."
+            )
+
+        if not agent.handoff_router:
+            logger.info(
+                "No handoff router defined for agent", agent=current_agent_name
+            )
+            span.add_event("no_router")
+            return None  # Indicate no handoff
+
+        logger.debug(
+            f"Using router {agent.handoff_router.__class__.__name__}",
+            agent=agent.name,
+        )
+        try:
+            # Execute the routing logic
+            handoff_data: (
+                HandOffRequest | Callable
+            ) = await agent.handoff_router.route(agent, result, context)
+
+            # Handle callable handoff functions - This is complex in distributed systems.
+            # Consider if this pattern should be supported or if routing should always
+            # return serializable data directly. Executing arbitrary code from context
+            # within an activity can have side effects and security implications.
+            # Assuming for now it MUST return HandOffRequest or structure convertible to it.
+            if callable(handoff_data):
+                logger.warning(
+                    "Callable handoff detected - executing function.",
+                    agent=agent.name,
+                )
+                # Ensure context is available if the callable needs it
+                try:
+                    handoff_data = handoff_data(
+                        context, result
+                    )  # Potential side effects
+                    if not isinstance(handoff_data, HandOffRequest):
+                        logger.error(
+                            "Handoff function did not return a HandOffRequest object.",
+                            agent=agent.name,
+                        )
+                        raise TypeError(
+                            "Handoff function must return a HandOffRequest object."
+                        )
+                except Exception as e:
+                    logger.error(
+                        "Handoff function execution failed",
+                        agent=agent.name,
+                        error=str(e),
+                        exc_info=True,
+                    )
+                    span.record_exception(e)
+                    raise  # Propagate error
+
+            # Ensure we have a HandOffRequest object after potentially calling function
+            if not isinstance(handoff_data, HandOffRequest):
+                logger.error(
+                    "Router returned unexpected type",
+                    type=type(handoff_data).__name__,
+                    agent=agent.name,
+                )
+                raise TypeError(
+                    f"Router for agent '{agent.name}' did not return a HandOffRequest object."
+                )
+
+            # Ensure agent instance is converted to name for serialization across boundaries
+            if isinstance(handoff_data.next_agent, FlockAgent):
+                handoff_data.next_agent = handoff_data.next_agent.name
+
+            # If router logic determines no further agent, return None
+            if not handoff_data.next_agent:
+                logger.info("Router determined no next agent", agent=agent.name)
+                span.add_event("no_next_agent_from_router")
+                return None
+
+            logger.info(
+                "Handoff determined",
+                next_agent=handoff_data.next_agent,
+                agent=agent.name,
+            )
+            span.set_attribute("next_agent", handoff_data.next_agent)
+            # Return the serializable HandOffRequest data using Pydantic's export method
+            return handoff_data.model_dump(
+                mode="json"
+            )  # Ensure JSON-serializable
+
+        except Exception as e:
+            # Catch potential errors during routing execution
+            logger.error(
+                "Router execution failed",
+                agent=agent.name,
+                error=str(e),
+                exc_info=True,
+            )
+            span.record_exception(e)
+            # Let Temporal handle the activity failure based on retry policy
+            raise

flock/workflow/flock_workflow.py

@@ -1,12 +1,22 @@
 from datetime import timedelta
+from typing import Any
 
 from temporalio import workflow
 
-from flock.core.context.context import FlockContext
-from flock.core.logging.logging import get_logger
-from flock.workflow.activities import run_agent
-
-# Import activity, passing it through the sandbox without reloading the module
+# Import activities from the new file
+with workflow.unsafe.imports_passed_through():
+    from flock.core.context.context import AgentDefinition, FlockContext
+    from flock.core.context.context_vars import FLOCK_CURRENT_AGENT
+    from flock.core.flock_router import HandOffRequest
+    from flock.core.logging.logging import get_logger
+    from flock.workflow.agent_execution_activity import (
+        determine_next_agent,
+        execute_single_agent,
+    )
+    from flock.workflow.temporal_config import (
+        TemporalActivityConfig,
+        TemporalRetryPolicyConfig,
+    )
 
 
 logger = get_logger("workflow")
@@ -14,45 +24,202 @@ logger = get_logger("workflow")
 
 @workflow.defn
 class FlockWorkflow:
-    def __init__(self) -> None:
-        self.context = None
+    # No need for __init__ storing context anymore if passed to run
 
     @workflow.run
-    async def run(self, context_dict: dict) -> dict:
-        self.context = FlockContext.from_dict(context_dict)
-        self.context.workflow_id = workflow.info().workflow_id
-        self.context.workflow_timestamp = workflow.info().start_time.strftime("%Y-%m-%d %H:%M:%S")
+    async def run(self, workflow_args: dict[str, Any]) -> dict:
+        # --- Workflow Initialization ---
+        # Arguments are packed into a single dictionary
+        context_dict = workflow_args["context_dict"]
+        default_retry_config_dict = workflow_args["default_retry_config_dict"]
+
+        # Deserialize context and default retry config
+        context = FlockContext.from_dict(context_dict)
+        default_retry_config = TemporalRetryPolicyConfig.model_validate(
+            default_retry_config_dict
+        )
+
+        context.workflow_id = workflow.info().workflow_id
+        context.workflow_timestamp = workflow.info().start_time.strftime(
+            "%Y-%m-%d %H:%M:%S"
+        )
+
+        current_agent_name = context.get_variable(FLOCK_CURRENT_AGENT)
+        final_result = None
+        previous_agent_name = (
+            None  # Keep track of the agent that called the current one
+        )
+
+        logger.info(
+            "Starting workflow execution",
+            workflow_id=context.workflow_id,
+            start_time=context.workflow_timestamp,
+            initial_agent=current_agent_name,
+        )
 
         try:
-            logger.info(
-                "Starting workflow execution",
-                timestamp=self.context.workflow_timestamp,
-            )
+            while current_agent_name:
+                logger.info(
+                    "Executing agent activity", agent=current_agent_name
+                )
 
-            result = await workflow.execute_activity(
-                run_agent,
-                self.context,
-                start_to_close_timeout=timedelta(minutes=5),
-            )
+                # --- Determine Activity Settings ---
+                agent_def: AgentDefinition | None = (
+                    context.get_agent_definition(current_agent_name)
+                )
+                agent_activity_config: TemporalActivityConfig | None = None
+                final_retry_config = (
+                    default_retry_config  # Start with the workflow default
+                )
+
+                if agent_def and agent_def.agent_data.get(
+                    "temporal_activity_config"
+                ):
+                    try:
+                        agent_activity_config = (
+                            TemporalActivityConfig.model_validate(
+                                agent_def.agent_data["temporal_activity_config"]
+                            )
+                        )
+                        logger.debug(
+                            f"Loaded agent-specific temporal config for {current_agent_name}"
+                        )
+                    except Exception as e:
+                        logger.warn(
+                            f"Failed to validate agent temporal config for {current_agent_name}: {e}. Using defaults."
+                        )
+
+                # Layering logic: Agent config overrides workflow default config
+                activity_task_queue = (
+                    workflow.info().task_queue
+                )  # Default to workflow task queue
+                activity_timeout = timedelta(
+                    minutes=5
+                )  # Fallback default timeout
+
+                if agent_activity_config:
+                    activity_task_queue = (
+                        agent_activity_config.task_queue or activity_task_queue
+                    )
+                    activity_timeout = (
+                        agent_activity_config.start_to_close_timeout
+                        or activity_timeout
+                    )
+                    if agent_activity_config.retry_policy:
+                        final_retry_config = agent_activity_config.retry_policy
+
+                # Convert config to actual Temporal object
+                final_retry_policy = final_retry_config.to_temporalio_policy()
+
+                logger.debug(
+                    f"Final activity settings for {current_agent_name}: "
+                    f"queue='{activity_task_queue}', timeout={activity_timeout}, "
+                    f"retries={final_retry_policy.maximum_attempts}"
+                )
 
-            self.context.set_variable(
+                # --- Execute the current agent activity ---
+                agent_result = await workflow.execute_activity(
+                    execute_single_agent,
+                    args=[current_agent_name, context],
+                    task_queue=activity_task_queue,  # Use determined task queue
+                    start_to_close_timeout=activity_timeout,  # Use determined timeout
+                    retry_policy=final_retry_policy,  # Use determined retry policy
+                )
+
+                # Record the execution in the context history
+                # Note: The 'called_from' is the agent *before* this one
+                context.record(
+                    agent_name=current_agent_name,
+                    data=agent_result,
+                    timestamp=workflow.now().isoformat(),  # Use deterministic workflow time
+                    hand_off=None,  # Will be updated if handoff occurs
+                    called_from=previous_agent_name,  # Pass the correct previous agent
+                )
+
+                final_result = agent_result  # Store the result of the last successful agent
+
+                logger.info(
+                    "Determining next agent activity",
+                    current_agent=current_agent_name,
+                )
+                # --- Determine the next agent activity (using workflow defaults for now) ---
+                # We could apply similar config logic to determine_next_agent if needed
+                handoff_data_dict = await workflow.execute_activity(
+                    determine_next_agent,
+                    args=[current_agent_name, agent_result, context],
+                    # Using sensible defaults, but could be configured via workflow_config?
+                    start_to_close_timeout=timedelta(minutes=1),
+                    retry_policy=default_retry_config.to_temporalio_policy(),  # Use default retry
+                )
+
+                # Update previous agent name for the next loop iteration
+                previous_agent_name = current_agent_name
+
+                if handoff_data_dict:
+                    logger.debug(
+                        "Handoff data received", data=handoff_data_dict
+                    )
+                    # Deserialize handoff data back into Pydantic model for easier access
+                    handoff_request = HandOffRequest.model_validate(
+                        handoff_data_dict
+                    )
+
+                    # Update context based on handoff overrides
+                    if handoff_request.override_context:
+                        context.state.update(handoff_request.override_context)
+                        logger.info("Context updated based on handoff override")
+
+                    # Update the last record's handoff information
+                    if context.history:
+                        context.history[-1].hand_off = handoff_data_dict
+
+                    # Set the next agent
+                    current_agent_name = handoff_request.next_agent
+                    if current_agent_name:
+                        context.set_variable(
+                            FLOCK_CURRENT_AGENT, current_agent_name
+                        )
+                        logger.info("Next agent set", agent=current_agent_name)
+                    else:
+                        logger.info(
+                            "Handoff requested termination (no next agent)"
+                        )
+                        break  # Exit loop if router explicitly returned no next agent
+
+                else:
+                    # No handoff data returned (no router or router returned None)
+                    logger.info("No handoff occurred, workflow terminating.")
+                    current_agent_name = None  # End the loop
+
+            # --- Workflow Completion ---
+            logger.success(
+                "Workflow completed successfully",
+                final_agent=previous_agent_name,
+            )
+            context.set_variable(
                 "flock.result",
                 {
-                    "result": result,
+                    "result": final_result,  # Return the last agent's result
                     "success": True,
                 },
             )
-
-            logger.success("Workflow completed successfully")
-            return result
+            return final_result  # Return the actual result of the last agent
 
         except Exception as e:
+            # Catch exceptions from activities (e.g., after retries fail)
+            # or workflow logic errors
             logger.exception("Workflow execution failed", error=str(e))
-            self.context.set_variable(
+            context.set_variable(
                 "flock.result",
                 {
-                    "result": f"Failed: {e}",
+                    "result": f"Workflow failed: {e}",
                     "success": False,
                 },
             )
-            return self.context
+            # It's often better to let Temporal record the failure status
+            # by re-raising the exception rather than returning a custom error dict.
+            # However, returning the context might be useful for debugging.
+            # Consider re-raising: raise
+            return context.model_dump(
+                mode="json"
+            )  # Return context state on failure

flock/workflow/temporal_config.py

@@ -0,0 +1,96 @@
+# src/flock/config/temporal_config.py
+
+"""Pydantic models for configuring Temporal execution settings."""
+
+from __future__ import annotations
+
+from datetime import timedelta
+from typing import TYPE_CHECKING
+
+# Conditionally import for type hinting only
+if TYPE_CHECKING:
+    from temporalio.common import RetryPolicy
+
+# Note: Importing temporalio types directly into config models can complicate serialization
+# if these models are meant to be purely data containers (e.g., for YAML/JSON).
+# We define the structure and provide a helper method to convert to the actual Temporal object.
+# Be careful if using workflow/activity decorators directly on methods within these config models.
+from pydantic import BaseModel, Field
+
+
+class TemporalRetryPolicyConfig(BaseModel):
+    """Configuration parameters for Temporal Retry Policies."""
+
+    initial_interval: timedelta = Field(
+        default=timedelta(seconds=1),
+        description="Initial delay before the first retry.",
+    )
+    backoff_coefficient: float = Field(
+        default=2.0, description="Multiplier for the delay between retries."
+    )
+    maximum_interval: timedelta | None = Field(
+        default=timedelta(seconds=100),
+        description="Maximum delay between retries.",
+    )
+    maximum_attempts: int = Field(
+        default=3,
+        description="Maximum number of retry attempts (0 means no retries after first failure).",
+    )
+    non_retryable_error_types: list[str] = Field(
+        default_factory=list,
+        description="List of error type names (strings) that should not be retried.",
+    )
+
+    # Helper to convert to actual Temporalio object when needed (e.g., in workflow/executor)
+    def to_temporalio_policy(self) -> RetryPolicy:
+        # Import locally to avoid making temporalio a hard dependency of the config module itself
+        # The type hint RetryPolicy is now available due to TYPE_CHECKING block
+        from temporalio.common import RetryPolicy
+
+        return RetryPolicy(
+            initial_interval=self.initial_interval,
+            backoff_coefficient=self.backoff_coefficient,
+            maximum_interval=self.maximum_interval,
+            maximum_attempts=self.maximum_attempts,
+            non_retryable_error_types=self.non_retryable_error_types,
+        )
+
+
+class TemporalWorkflowConfig(BaseModel):
+    """Configuration specific to Temporal Workflow Execution for a Flock."""
+
+    task_queue: str = Field(
+        default="flock-queue",
+        description="Default task queue for the workflow execution.",
+    )
+    workflow_execution_timeout: timedelta | None = Field(
+        default=None,  # Default to no timeout (Temporal server default)
+        description="Total time limit for the workflow execution.",
+    )
+    workflow_run_timeout: timedelta | None = Field(
+        default=None,  # Default to no timeout (Temporal server default)
+        description="Time limit for a single workflow run attempt.",
+    )
+    # Default retry policy for activities if not specified per-agent
+    default_activity_retry_policy: TemporalRetryPolicyConfig = Field(
+        default_factory=TemporalRetryPolicyConfig,
+        description="Default retry policy applied to activities if not overridden by the agent.",
+    )
+
+
+class TemporalActivityConfig(BaseModel):
+    """Configuration specific to Temporal Activity Execution (per Agent)."""
+
+    task_queue: str | None = Field(
+        default=None,
+        description="Specific task queue for this agent's activity execution (overrides workflow default).",
+    )
+    start_to_close_timeout: timedelta | None = Field(
+        default=timedelta(minutes=5),  # Default to 5 minutes
+        description="Time limit for a single activity attempt.",
+    )
+    retry_policy: TemporalRetryPolicyConfig | None = Field(
+        default=None,
+        description="Specific retry policy for this activity (overrides workflow default).",
+    )
+    # Other timeouts like schedule_to_start, heartbeat_timeout could be added here if needed

flock/workflow/temporal_setup.py

@@ -1,5 +1,3 @@
-import asyncio
-import traceback
 import uuid
 
 from temporalio.client import Client
@@ -7,34 +5,33 @@ from temporalio.worker import Worker
 
 
 async def create_temporal_client() -> Client:
+    # Consider making the address configurable
     client = await Client.connect("localhost:7233")
     return client
 
 
-async def setup_worker(workflow, activity) -> Client:
-    """Setup the worker and return the result of the worker.run() call."""
-    try:
-        # debug
-        # print("workflow =", workflow)
-        # print("isclass     :", inspect.isclass(workflow))
-        # print(
-        #     "has _defn   :",
-        #     hasattr(workflow, "_temporalio_workflow_definition"),
-        # )
-        worker_client = await create_temporal_client()
-        worker = Worker(
-            worker_client,
-            task_queue="flock-queue",
-            workflows=[workflow],
-            activities=[activity],
-        )
-        result = asyncio.create_task(worker.run())
-        await asyncio.sleep(1)
-        return result
-    except Exception:
-        print("\n=== Worker construction failed ===")
-        traceback.print_exc()
-        raise
+async def setup_worker(
+    client: Client, task_queue: str, workflow: type, activities: list
+) -> Worker:
+    """Creates and configures a worker instance, but does not run it.
+
+    Args:
+        client: The Temporal client to associate with the worker.
+        task_queue: The task queue the worker should listen on.
+        workflow: The workflow class definition.
+        activities: A list of activity functions.
+
+    Returns:
+        A configured Worker instance.
+    """
+    # Creates and configures the worker instance
+    worker = Worker(
+        client,
+        task_queue=task_queue,
+        workflows=[workflow],
+        activities=activities,
+    )
+    return worker  # Return the configured worker instance
 
 
 async def run_worker(client: Client, task_queue: str, workflows, activities):