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

flock/core/context/context.py

@@ -20,7 +20,7 @@ class AgentRunRecord(BaseModel):
     data: dict[str, Any] = Field(default_factory=dict)
     timestamp: str = Field(default="")
     hand_off: dict | None = Field(default_factory=dict)
-    called_from: str = Field(default="")
+    called_from: str | None = Field(default=None)
 
 
 class AgentDefinition(BaseModel):
@@ -132,6 +132,15 @@ class FlockContext(Serializable, BaseModel):
     def get_agent_definition(self, agent_name: str) -> AgentDefinition | None:
         return self.agent_definitions.get(agent_name)
 
+    def get_last_agent_name(self) -> str | None:
+        """Returns the name of the agent from the most recent history record."""
+        if not self.history:
+            return None
+        last_record = self.history[-1]
+        # The 'called_from' field in the *next* record is the previous agent.
+        # However, to get the name of the *last executed agent*, we look at the 'agent' field.
+        return last_record.agent
+
     def add_agent_definition(
         self, agent_type: type, agent_name: str, agent_data: Any
     ) -> None:

flock/core/execution/temporal_executor.py

@@ -1,10 +1,23 @@
 # src/your_package/core/execution/temporal_executor.py
 
+import asyncio  # Import asyncio
+from typing import TYPE_CHECKING, Any
+
+from temporalio.worker import Worker  # Import Worker
+
+if TYPE_CHECKING:
+    from flock.core.flock import Flock  # Import Flock for type hinting
+
 from flock.core.context.context import FlockContext
 from flock.core.context.context_vars import FLOCK_RUN_ID
 from flock.core.logging.logging import get_logger
-from flock.workflow.activities import (
-    run_agent,  # Activity function used in Temporal
+from flock.workflow.agent_execution_activity import (
+    determine_next_agent,
+    execute_single_agent,
+)
+from flock.workflow.temporal_config import (
+    TemporalRetryPolicyConfig,
+    TemporalWorkflowConfig,
 )
 from flock.workflow.temporal_setup import create_temporal_client, setup_worker
 
@@ -12,14 +25,18 @@ logger = get_logger("flock")
 
 
 async def run_temporal_workflow(
+    flock_instance: "Flock",  # Accept Flock instance
     context: FlockContext,
     box_result: bool = True,
+    memo: dict[str, Any] | None = None,  # Add memo argument
 ) -> dict:
     """Execute the agent workflow via Temporal for robust, distributed processing.
 
     Args:
+        flock_instance: The Flock instance.
         context: The FlockContext instance with state and history.
         box_result: If True, wraps the result in a Box for nicer display.
+        memo: Optional dictionary of metadata to attach to the Temporal workflow.
 
     Returns:
         A dictionary containing the workflow result.
@@ -29,29 +46,121 @@ async def run_temporal_workflow(
             FlockWorkflow,  # Your workflow class
         )
 
-        logger.info("Setting up Temporal workflow")
-        await setup_worker(workflow=FlockWorkflow, activity=run_agent)
+        # Get workflow config from Flock instance or use defaults
+        wf_config = flock_instance.temporal_config or TemporalWorkflowConfig()
+
         logger.debug("Creating Temporal client")
         flock_client = await create_temporal_client()
-        workflow_id = context.get_variable(FLOCK_RUN_ID)
-        logger.info("Executing Temporal workflow", workflow_id=workflow_id)
-        context_dict = context.model_dump()
-        result = await flock_client.execute_workflow(
-            FlockWorkflow.run,
-            context_dict,
-            id=workflow_id,
-            task_queue="flock-queue",
-        )
 
-        agent_name = context.get_variable("FLOCK_CURRENT_AGENT")
-        logger.debug("Formatting Temporal result", agent=agent_name)
+        # Determine if we need to manage an in-process worker
+        start_worker_locally = flock_instance.temporal_start_in_process_worker
+
+        # Setup worker instance
+        worker: Worker | None = None
+        worker_task: asyncio.Task | None = None
+
+        if start_worker_locally:
+            logger.info(
+                f"Setting up temporary in-process worker for task queue '{wf_config.task_queue}'"
+            )
+            worker = await setup_worker(
+                flock_client,  # Pass the client
+                wf_config.task_queue,  # Pass the task queue
+                FlockWorkflow,
+                [execute_single_agent, determine_next_agent],
+            )
+
+            # Run the worker in the background
+            worker_task = asyncio.create_task(worker.run())
+            logger.info("Temporal worker started in background.")
+
+            # Allow worker time to start polling (heuristic for local testing)
+            await asyncio.sleep(2)
+        else:
+            logger.info(
+                "Skipping in-process worker startup. Assuming dedicated workers are running."
+            )
+
+        try:
+            workflow_id = context.get_variable(FLOCK_RUN_ID)
+            logger.info(
+                "Executing Temporal workflow",
+                workflow_id=workflow_id,
+                task_queue=wf_config.task_queue,
+            )
+
+            # Prepare the single workflow argument dictionary
+            workflow_args_dict = {
+                "context_dict": context.model_dump(mode="json"),
+                "default_retry_config_dict": (
+                    wf_config.default_activity_retry_policy.model_dump(
+                        mode="json"
+                    )
+                    if wf_config.default_activity_retry_policy
+                    else TemporalRetryPolicyConfig().model_dump(mode="json")
+                ),
+            }
+
+            # Start the workflow using start_workflow
+            handle = await flock_client.start_workflow(
+                FlockWorkflow.run,
+                # Pass the single dictionary as the only element in the args list
+                args=[workflow_args_dict],
+                id=workflow_id,
+                task_queue=wf_config.task_queue,
+                # Corrected timeout argument names
+                execution_timeout=wf_config.workflow_execution_timeout,
+                run_timeout=wf_config.workflow_run_timeout,
+                memo=memo or {},  # Pass memo if provided
+            )
+
+            logger.info(
+                "Workflow started, awaiting result...", workflow_id=handle.id
+            )
+            # Await the result from the handle
+            result = await handle.result()
+            logger.info("Workflow result received.")
+
+            agent_name = context.get_variable("FLOCK_CURRENT_AGENT")
+            logger.debug("Formatting Temporal result", agent=agent_name)
 
-        if box_result:
-            from box import Box
+            if box_result:
+                from box import Box
 
-            logger.debug("Boxing Temporal result")
-            return Box(result)
-        return result
+                logger.debug("Boxing Temporal result")
+                return Box(result)
+            return result
+        except Exception as e:
+            logger.error(
+                "Error during Temporal workflow execution or result retrieval",
+                error=e,
+            )
+            raise e  # Re-raise the exception after logging
+        finally:
+            # Ensure worker is shut down regardless of success or failure
+            if (
+                start_worker_locally
+                and worker
+                and worker_task
+                and not worker_task.done()
+            ):
+                logger.info("Shutting down temporal worker...")
+                await worker.shutdown()  # Await the shutdown coroutine
+                try:
+                    await asyncio.wait_for(
+                        worker_task, timeout=10.0
+                    )  # Wait for task to finish
+                    logger.info("Temporal worker shut down gracefully.")
+                except asyncio.TimeoutError:
+                    logger.warning(
+                        "Temporal worker shutdown timed out. Cancelling task."
+                    )
+                    worker_task.cancel()
+                except Exception as shutdown_err:
+                    logger.error(
+                        f"Error during worker shutdown: {shutdown_err}",
+                        exc_info=True,
+                    )
     except Exception as e:
         logger.error("Error executing Temporal workflow", error=e)
         raise e

flock/core/flock.py

@@ -39,6 +39,7 @@ from flock.core.flock_evaluator import FlockEvaluator
 from flock.core.logging.logging import LOGGERS, get_logger, get_module_loggers
 from flock.core.serialization.serializable import Serializable
 from flock.core.util.cli_helper import init_console
+from flock.workflow.temporal_config import TemporalWorkflowConfig
 
 # Import FlockAgent using TYPE_CHECKING to avoid circular import at runtime
 if TYPE_CHECKING:
@@ -102,6 +103,16 @@ class Flock(BaseModel, Serializable):
         default=True,
         description="If True, show the Flock banner on console interactions.",
     )
+    # --- Temporal Configuration (Optional) ---
+    temporal_config: TemporalWorkflowConfig | None = Field(
+        default=None,
+        description="Optional Temporal settings specific to the workflow execution for this Flock.",
+    )
+    # --- Temporal Dev/Test Setting ---
+    temporal_start_in_process_worker: bool = Field(
+        default=True,
+        description="If True (default) and enable_temporal=True, start a temporary in-process worker for development/testing convenience. Set to False when using dedicated workers.",
+    )
     # Internal agent storage - not part of the Pydantic model for direct serialization
     _agents: dict[str, FlockAgent]
     _start_agent_name: str | None = None  # For potential pre-configuration
@@ -122,6 +133,8 @@ class Flock(BaseModel, Serializable):
         enable_temporal: bool = False,
         enable_logging: bool | list[str] = False,
         agents: list[FlockAgent] | None = None,
+        temporal_config: TemporalWorkflowConfig | None = None,
+        temporal_start_in_process_worker: bool = True,
         **kwargs,
     ):
         """Initialize the Flock orchestrator."""
@@ -136,6 +149,8 @@ class Flock(BaseModel, Serializable):
             enable_temporal=enable_temporal,
             enable_logging=enable_logging,
             show_flock_banner=show_flock_banner,
+            temporal_config=temporal_config,
+            temporal_start_in_process_worker=temporal_start_in_process_worker,
             **kwargs,
         )
 
@@ -311,6 +326,7 @@ class Flock(BaseModel, Serializable):
         run_id: str = "",
         box_result: bool = True,
         agents: list[FlockAgent] | 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
@@ -348,7 +364,17 @@ class Flock(BaseModel, Serializable):
 
             # Check if start_agent is in agents
             if start_agent_name not in self._agents:
-                raise ValueError(f"Start agent '{start_agent_name}' not found.")
+                # Try loading from registry if not found locally yet
+                reg_agent = FlockRegistry.get_agent(start_agent_name)
+                if reg_agent:
+                    self.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."
+                    )
 
             run_input = input if input is not None else self._start_input
             effective_run_id = run_id or f"flockrun_{uuid.uuid4().hex[:8]}"
@@ -394,6 +420,15 @@ class Flock(BaseModel, Serializable):
                         agent_data=agent_dict_repr,  # Pass the serialized dict
                     )
 
+                # Add temporal config to context if enabled
+                if self.enable_temporal and self.temporal_config:
+                    # Store the workflow config dict for the executor/workflow to use
+                    # Using a specific key to avoid potential clashes in state
+                    run_context.set_variable(
+                        "flock.temporal_workflow_config",
+                        self.temporal_config.model_dump(mode="json"),
+                    )
+
                 logger.info(
                     "Starting agent execution",
                     agent=start_agent_name,
@@ -406,8 +441,15 @@ class Flock(BaseModel, Serializable):
                         run_context, box_result=False
                     )
                 else:
+                    # Pass the Flock instance itself to the executor
+                    # so it can access the temporal_config directly if needed
+                    # This avoids putting potentially large/complex config objects
+                    # directly into the context state that gets passed around.
                     result = await run_temporal_workflow(
-                        run_context, box_result=False
+                        self,  # Pass the Flock instance
+                        run_context,
+                        box_result=False,
+                        memo=memo,
                     )
 
                 span.set_attribute("result.type", str(type(result)))
@@ -668,6 +710,7 @@ class Flock(BaseModel, Serializable):
         # Import locally to prevent circular imports at module level if structure is complex
         from flock.core.serialization.flock_serializer import FlockSerializer
 
+        # Assuming FlockSerializer handles the nested temporal_config serialization
         return FlockSerializer.serialize(self, path_type=path_type)
 
     @classmethod
@@ -676,6 +719,7 @@ class Flock(BaseModel, Serializable):
         # Import locally
         from flock.core.serialization.flock_serializer import FlockSerializer
 
+        # Assuming FlockSerializer handles the nested temporal_config deserialization
         return FlockSerializer.deserialize(cls, data)
 
     # --- Static Method Loader (Delegates to loader module) ---