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

flock/workflow/activities.py

@@ -1,115 +1,175 @@
-"""Defines Temporal activities for running a chain of agents."""
-
-from datetime import datetime
-
-from temporalio import activity
-
-from flock.core.agent import Agent
-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
-from flock.core.handoff.handoff_base import HandoffBase
-from flock.core.logging import flock_logger
-
-
-@activity.defn
-async def run_agent(context: FlockContext) -> dict:
-    """Runs a chain of agents using the provided context.
-
-    The context contains:
-      - A state (e.g., "init_input", "current_agent", etc.),
-      - A history of agent runs.
-
-    Each agent uses the current value of the variable specified in its `input` attribute.
-    After each run, its output is merged into the context state.
-    The default handoff behavior is to fetch the next agent's input automatically from the context.
-    """
-    registry = Registry()
-    current_agent_name = context.get_variable(FLOCK_CURRENT_AGENT)
-    flock_logger.info(f"Starting agent chain with initial agent", agent=current_agent_name)
-
-    agent = registry.get_agent(current_agent_name)
-    if not agent:
-        flock_logger.error(f"Agent not found", agent=current_agent_name)
-        return {"error": f"Agent '{current_agent_name}' not found."}
-
-    while agent:
-        # Determine the input for this agent.
-        # (Preferably, the agent's input key is used; otherwise "init_input" is assumed.)
-        agent_input = (
-            context.get_variable(agent.input)
-            if getattr(agent, "input", None)
-            else context.get_variable(FLOCK_INITIAL_INPUT)
-        )
-        flock_logger.info(
-            "Determined agent input", agent=agent.name, input_source=getattr(agent, "input", "FLOCK_INITIAL_INPUT")
-        )
-
-        # Prepare a deep copy of the state for the agent's run.
-        local_context = FlockContext.from_dict(context.to_dict())
-        local_context.set_variable(
-            f"{agent.name}.{agent.input}", agent_input
-        )  # ensure the agent's expected input is set
-
-        # Execute the agent.
-        flock_logger.info("Executing agent", agent=agent.name)
-        result = await agent.run(local_context)
-        flock_logger.info("Agent execution completed", agent=agent.name)
-
-        # If no handoff is defined, return the result.
-        if not agent.hand_off:
-            # Record the agent's execution.
-            context.record(agent.name, result, timestamp=datetime.now(), hand_off=None)
-            flock_logger.info("No handoff defined, completing chain", agent=agent.name)
-            return result
-
-        # Determine the next agent.
-        if callable(agent.hand_off):
-            # The handoff function may override the default behavior:
-            # it can explicitly return a dict with "next_agent" and optionally "input"
-            flock_logger.info("Executing handoff function", agent=agent.name)
-            handoff_data: HandoffBase = agent.hand_off(context, result)
-            if isinstance(handoff_data.next_agent, Agent):
-                next_agent_name = handoff_data.next_agent.name
-            else:
-                next_agent_name = handoff_data.next_agent
-
-            # Use the provided new input if present, otherwise let the context update be automatic.
-            if handoff_data.input:
-                context.state["init_input"] = handoff_data["input"]
-                flock_logger.debug("Using handoff-provided input", agent=agent.name)
-            if handoff_data.context_params:
-                context.state.update(handoff_data.context_params)
-                flock_logger.debug("Updated context with handoff params", agent=agent.name)
-        elif isinstance(agent.hand_off, str | Agent):
-            next_agent_name = agent.hand_off if isinstance(agent.hand_off, str) else agent.hand_off.name
-
-        else:
-            return {"error": "Unsupported hand_off type."}
-        context.record(agent.name, result, timestamp=datetime.now(), hand_off=next_agent_name)
-        # Update the current agent and prepare the next input automatically.
-        next_agent = registry.get_agent(next_agent_name)
-        if not next_agent:
-            flock_logger.error("Next agent not found", agent=next_agent_name)
-            return {"error": f"Next agent '{next_agent_name}' not found."}
-
-        context.state["current_agent"] = next_agent.name
-        # By default, the next input is determined from the current state using next_input_for().
-        context.state["init_input"] = context.next_input_for(next_agent)
-        flock_logger.info("Handing off to next agent", current=agent.name, next=next_agent.name)
-        agent = next_agent
-
-    # If the loop ever terminates without a final result, return the latest state value.
-    flock_logger.info("Chain completed, returning final state")
-    return context.get_variable("init_input")
-
-
-@activity.defn
-async def get_next_agent(name: str) -> Agent | None:
-    """Retrieves the agent with the given name from the registry."""
-    flock_logger.debug("Looking up next agent", agent=name)
-    registry = Registry()
-    agent = registry.get_agent(name)
-    if not agent:
-        flock_logger.warning("Next agent not found", agent=name)
-    return agent
+"""Defines Temporal activities for running a chain of agents with logging and tracing."""
+
+from datetime import datetime
+
+from opentelemetry import trace
+from temporalio import activity
+
+from flock.core.context.context import FlockContext
+from flock.core.context.context_vars import FLOCK_CURRENT_AGENT
+from flock.core.flock_agent import FlockAgent, HandOff
+from flock.core.logging.formatters.base_formatter import FormatterOptions
+from flock.core.logging.formatters.formatter_factory import FormatterFactory
+from flock.core.logging.logging import get_logger
+from flock.core.registry.agent_registry import Registry
+from flock.core.util.input_resolver import resolve_inputs
+
+logger = get_logger("activities")
+tracer = trace.get_tracer(__name__)
+
+
+@activity.defn
+async def run_agent(
+    context: FlockContext, output_formatter: FormatterOptions = None
+) -> dict:
+    """Runs a chain of agents using the provided context.
+
+    The context contains state, history, and agent definitions.
+    After each agent run, its output is merged into the context.
+    """
+    # Start a top-level span for the entire run_agent activity.
+    with tracer.start_as_current_span("run_agent") as span:
+        registry = Registry()
+        previous_agent_name = ""
+        if isinstance(context, dict):
+            context = FlockContext.from_dict(context)
+        current_agent_name = context.get_variable(FLOCK_CURRENT_AGENT)
+        span.set_attribute("initial.agent", current_agent_name)
+        logger.info("Starting agent chain", initial_agent=current_agent_name)
+
+        agent = registry.get_agent(current_agent_name)
+        agent.resolve_callables(context=context)
+        if not agent:
+            logger.error("Agent not found", agent=current_agent_name)
+            span.record_exception(
+                Exception(f"Agent '{current_agent_name}' not found")
+            )
+            return {"error": f"Agent '{current_agent_name}' not found."}
+
+        # Loop over agents in the chain.
+        while agent:
+            # Create a nested span for this iteration.
+            with tracer.start_as_current_span("agent_iteration") as iter_span:
+                iter_span.set_attribute("agent.name", agent.name)
+
+                # Resolve inputs for the agent.
+                agent_inputs = resolve_inputs(
+                    agent.input, context, previous_agent_name
+                )
+                iter_span.add_event(
+                    "resolved inputs", attributes={"inputs": str(agent_inputs)}
+                )
+
+                # Execute the agent with its own span.
+                with tracer.start_as_current_span("execute_agent") as exec_span:
+                    logger.info("Executing agent", agent=agent.name)
+                    try:
+                        result = await agent.run(agent_inputs)
+                        exec_span.set_attribute("result", str(result))
+                        logger.debug(
+                            "Agent execution completed", agent=agent.name
+                        )
+                    except Exception as e:
+                        logger.error(
+                            "Agent execution failed",
+                            agent=agent.name,
+                            error=str(e),
+                        )
+                        exec_span.record_exception(e)
+                        raise
+
+                # Optionally display formatted output.
+                display_output = True
+                if agent.config:
+                    display_output = not agent.config.disable_output
+
+                if output_formatter and display_output:
+                    formatter = FormatterFactory.create_formatter(
+                        output_formatter
+                    )
+                    formatter.display(
+                        result, agent.name, output_formatter.wait_for_input
+                    )
+
+                # If there is no handoff, record the result and finish.
+                if not agent.hand_off:
+                    context.record(
+                        agent.name,
+                        result,
+                        timestamp=datetime.now(),
+                        hand_off=None,
+                        called_from=previous_agent_name,
+                    )
+                    logger.info(
+                        "No handoff defined, completing chain", agent=agent.name
+                    )
+                    iter_span.add_event("chain completed")
+                    return result
+
+                # Determine the next agent.
+                handoff_data = HandOff()
+                if callable(agent.hand_off):
+                    logger.debug("Executing handoff function", agent=agent.name)
+                    try:
+                        handoff_data = agent.hand_off(context, result)
+                        if isinstance(handoff_data.next_agent, FlockAgent):
+                            handoff_data.next_agent = (
+                                handoff_data.next_agent.name
+                            )
+                    except Exception as e:
+                        logger.error(
+                            "Handoff function error",
+                            agent=agent.name,
+                            error=str(e),
+                        )
+                        iter_span.record_exception(e)
+                        return {"error": f"Handoff function error: {e}"}
+                elif isinstance(agent.hand_off, str | FlockAgent):
+                    handoff_data.next_agent = (
+                        agent.hand_off
+                        if isinstance(agent.hand_off, str)
+                        else agent.hand_off.name
+                    )
+                else:
+                    logger.error("Unsupported hand_off type", agent=agent.name)
+                    iter_span.add_event("unsupported hand_off type")
+                    return {"error": "Unsupported hand_off type."}
+
+                # Record the agent run in the context.
+                context.record(
+                    agent.name,
+                    result,
+                    timestamp=datetime.now().isoformat(),
+                    hand_off=handoff_data,
+                    called_from=previous_agent_name,
+                )
+                previous_agent_name = agent.name
+
+                # Prepare the next agent.
+                try:
+                    agent = registry.get_agent(handoff_data.next_agent)
+                    agent.resolve_callables(context=context)
+                    if not agent:
+                        logger.error(
+                            "Next agent not found",
+                            agent=handoff_data.next_agent,
+                        )
+                        iter_span.record_exception(
+                            Exception(
+                                f"Next agent '{handoff_data.next_agent}' not found"
+                            )
+                        )
+                        return {
+                            "error": f"Next agent '{handoff_data.next_agent}' not found."
+                        }
+
+                    context.set_variable(FLOCK_CURRENT_AGENT, agent.name)
+                    logger.info("Handing off to next agent", next=agent.name)
+                    iter_span.set_attribute("next.agent", agent.name)
+                except Exception as e:
+                    logger.error("Error during handoff", error=str(e))
+                    iter_span.record_exception(e)
+                    return {"error": f"Error during handoff: {e}"}
+
+        # If the loop exits unexpectedly, return the initial input.
+        return context.get_variable("init_input")

flock/workflow/agent_activities.py

@@ -1,26 +1,24 @@
-from temporalio import activity
-
-from flock.agents.declarative_agent import DeclarativeAgent
-from flock.core.context import FlockContext
-
-
-@activity.defn
-async def run_declarative_agent_activity(params: dict) -> dict:
-    """Temporal activity to run a declarative (or batch) agent.
-
-    Expects a dictionary with:
-      - "agent_data": a dict representation of the agent (as produced by .dict()),
-      - "context_data": a dict containing the FlockContext state and optionally other fields.
-
-    The activity reconstructs the agent and a FlockContext, then calls the agent’s _evaluate() method.
-    """
-    agent_data = params.get("agent_data")
-    context_data = params.get("context_data", {})
-    # Reconstruct the agent from its serialized representation.
-    agent = DeclarativeAgent.parse_obj(agent_data)
-    # Reconstruct the FlockContext from the state.
-    state = context_data.get("state", {})
-    # (For simplicity, we ignore history and agent_definitions here. You can extend this if needed.)
-    context = FlockContext(state=state)
-    result = await agent._evaluate(context)
-    return result
+from temporalio import activity
+
+from flock.core.context.context import FlockContext
+from flock.core.flock_agent import FlockAgent
+
+
+@activity.defn
+async def run_declarative_agent_activity(params: dict) -> dict:
+    """Temporal activity to run a declarative (or batch) agent.
+
+    Expects a dictionary with:
+      - "agent_data": a dict representation of the agent (as produced by .dict()),
+      - "context_data": a dict containing the FlockContext state and optionally other fields.
+
+    The activity reconstructs the agent and a FlockContext, then calls the agent’s _evaluate() method.
+    """
+    agent_data = params.get("agent_data")
+    context_data = params.get("context_data", {})
+    # Reconstruct the agent from its serialized representation.
+    agent = FlockAgent.from_dict(agent_data)
+    # Reconstruct the FlockContext from the state.
+    context = FlockContext.from_dict(context_data)
+    result = await agent.evaluate(context)
+    return result

flock/workflow/temporal_setup.py

@@ -1,37 +1,38 @@
-import asyncio
-import uuid
-
-from temporalio.client import Client
-from temporalio.worker import Worker
-
-
-async def create_temporal_client() -> Client:
-    client = await Client.connect("localhost:7233")
-    return client
-
-async def setup_worker(workflow,activity) -> Client:
-    worker_client = await create_temporal_client()
-    worker = Worker(worker_client, task_queue="flock-queue", workflows=[workflow], activities=[activity])
-    asyncio.create_task(worker.run())
-    await asyncio.sleep(1)
-
-
-async def run_worker(client: Client, task_queue: str, workflows, activities):
-    worker = Worker(client, task_queue=task_queue, workflows=workflows, activities=activities)
-    await worker.run()
-
-
-async def run_activity(client: Client, name: str, func, param):
-    run_id = f"{name}_{uuid.uuid4().hex[:4]}"
-
-    try:
-        result = await client.execute_activity(
-            func,
-            param,
-            id=run_id,
-            task_queue="flock-queue",
-            start_to_close_timeout=300,  # e.g., 5 minutes
-        )
-        return result
-    except Exception:
-        raise
+import asyncio
+import uuid
+
+from temporalio.client import Client
+from temporalio.worker import Worker
+
+
+async def create_temporal_client() -> Client:
+    client = await Client.connect("localhost:7233")
+    return client
+
+
+async def setup_worker(workflow, activity) -> Client:
+    worker_client = await create_temporal_client()
+    worker = Worker(worker_client, task_queue="flock-queue", workflows=[workflow], activities=[activity])
+    asyncio.create_task(worker.run())
+    await asyncio.sleep(1)
+
+
+async def run_worker(client: Client, task_queue: str, workflows, activities):
+    worker = Worker(client, task_queue=task_queue, workflows=workflows, activities=activities)
+    await worker.run()
+
+
+async def run_activity(client: Client, name: str, func, param):
+    run_id = f"{name}_{uuid.uuid4().hex[:4]}"
+
+    try:
+        result = await client.execute_activity(
+            func,
+            param,
+            id=run_id,
+            task_queue="flock-queue",
+            start_to_close_timeout=300,  # e.g., 5 minutes
+        )
+        return result
+    except Exception:
+        raise

flock/workflow/workflow.py

@@ -1,53 +1,58 @@
-from datetime import timedelta
-
-from temporalio import workflow
-
-from flock.core.context import FlockContext
-from flock.core.logging import flock_logger
-from flock.workflow.activities import run_agent
-
-# Import activity, passing it through the sandbox without reloading the module
-with workflow.unsafe.imports_passed_through():
-    from flock.workflow.activities import run_agent
-
-
-@workflow.defn
-class FlockWorkflow:
-    def __init__(self) -> None:
-        self.context = None
-
-    @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")
-
-        try:
-            flock_logger.set_context(workflow_id=self.context.workflow_id)
-            flock_logger.workflow_event(f"Starting workflow execution at {self.context.workflow_timestamp}")
-
-            result = await workflow.execute_activity(
-                run_agent, self.context, start_to_close_timeout=timedelta(minutes=5)
-            )
-
-            self.context.set_variable(
-                "flock.result",
-                {
-                    "result": result,
-                    "success": True,
-                },
-            )
-
-            flock_logger.workflow_event("Workflow completed successfully")
-            return result
-
-        except Exception as e:
-            flock_logger.error(f"Workflow execution failed: {e}")
-            self.context.set_variable(
-                "flock.result",
-                {
-                    "result": f"Failed: {e}",
-                    "success": False,
-                },
-            )
-            return self.context
+from datetime import timedelta
+
+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
+
+
+logger = get_logger("workflow")
+
+
+@workflow.defn
+class FlockWorkflow:
+    def __init__(self) -> None:
+        self.context = None
+
+    @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")
+
+        try:
+            logger.info(
+                "Starting workflow execution",
+                timestamp=self.context.workflow_timestamp,
+            )
+
+            result = await workflow.execute_activity(
+                run_agent,
+                self.context,
+                start_to_close_timeout=timedelta(minutes=5),
+            )
+
+            self.context.set_variable(
+                "flock.result",
+                {
+                    "result": result,
+                    "success": True,
+                },
+            )
+
+            logger.success("Workflow completed successfully")
+            return result
+
+        except Exception as e:
+            logger.exception("Workflow execution failed", error=str(e))
+            self.context.set_variable(
+                "flock.result",
+                {
+                    "result": f"Failed: {e}",
+                    "success": False,
+                },
+            )
+            return self.context